package mr

import (
	"bufio"
	"fmt"
	"log"
	"net"
	"net/http"
	"net/rpc"
	"os"
	"path/filepath"
	"sync"
	"time"
)

type Task struct {
	State    int // 0：未开始  1：已经派发，执行中 2：执行完成
	WorkerID int //派发给的workerID
}

type WorkerState struct {
	TaskID           int
	lastHeatbeatTime time.Time
}

type Coordinator struct {
	mu               sync.Mutex
	nReduce          int
	mapFinished      bool
	Finished         bool
	WorkerID2TaskID  map[int]WorkerState
	WorkerID2TaskID1 map[int]WorkerState
	MapTaskList      []Task
	ReduceTaskList   []Task
	workerID         int
	files            []string
}

// Your code here -- RPC handlers for the worker to call.

// an example RPC handler.
//
// the RPC argument and reply types are defined in rpc.go.
func (c *Coordinator) Example(args *ExampleArgs, reply *ExampleReply) error {
	reply.Y = args.X + 1
	return nil
}

func (c *Coordinator) GenerateWorkerID() int {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.workerID++
	return c.workerID
}

func (c *Coordinator) GetTask(req *GetTaskParam, res *GetTaskRes) error {
	// if req.WorkerID <= 0 {
	// 	res.WorkerID = c.GenerateWorkerID()
	// } else {
	// 	res.WorkerID = req.WorkerID
	// }
	res.WorkerID = req.WorkerID
	c.mu.Lock()
	defer c.mu.Unlock()

	//节点完成任务
	// fmt.Printf("coordinator get %+v\n", req)

	if !c.mapFinished {
		if req.TaskState == 2 {
			tid := c.WorkerID2TaskID[res.WorkerID].TaskID
			c.MapTaskList[tid].State = 2
			delete(c.WorkerID2TaskID, res.WorkerID)
			//检查是否所有任务都已经完成
			mapFinished := true
			for _, task := range c.MapTaskList {
				if task.State != 2 {
					mapFinished = false
					// fmt.Printf("%d false\n", i)
					break
				}
			}
			if mapFinished {
				c.mapFinished = true
			}
			// fmt.Printf("mapFinished: %v\n", c.mapFinished)
		}
	} else {
		if req.TaskState == 2 {
			tid := c.WorkerID2TaskID1[res.WorkerID].TaskID
			c.ReduceTaskList[tid].State = 2
			delete(c.WorkerID2TaskID1, res.WorkerID)
			//检查是否所有任务都已经完成
			Finished := true
			for _, task := range c.ReduceTaskList {
				if task.State != 2 {
					Finished = false
					break
				}
			}
			if Finished {
				c.Finished = true
			}
		}
	}

	//正在干活 就不派活了
	if req.TaskState == 1 {
		res.TaskType = 0
		return nil
	}

	if c.Finished {
		res.TaskType = 0
		return nil
	}
	res.Files = c.files
	//mp任务未全部完成
	if !c.mapFinished {
		cnt := 0
		for _, task := range c.MapTaskList {
			if task.State != 0 {
				cnt++
			}
		}
		if cnt == len(c.MapTaskList) {
			res.TaskType = 3
			return nil
		}
		res.TaskType = 1
		res.NReduce = c.nReduce
		for i, task := range c.MapTaskList {
			if task.State == 0 {
				res.TaskID = i
				c.MapTaskList[i].State = 1
				c.WorkerID2TaskID[res.WorkerID] = WorkerState{TaskID: i, lastHeatbeatTime: time.Now()}
				break
			}
		}
	} else {
		res.TaskType = 2
		for i, task := range c.ReduceTaskList {
			if task.State == 0 {
				res.TaskID = i
				c.ReduceTaskList[i].State = 1
				c.WorkerID2TaskID1[res.WorkerID] = WorkerState{TaskID: i, lastHeatbeatTime: time.Now()}
				break
			}
		}

	}
	return nil
}

func (c *Coordinator) CallHeatbeatCheck(req *HeatbeatReq, resp *HeatbeatResp) error {
	id := req.WorkerID
	resp.CoordinatorState = 1
	c.mu.Lock()
	defer c.mu.Unlock()
	// fmt.Printf("heartbeat from %d \n", req.WorkerID)
	// fmt.Printf("map: %+v\n", c.WorkerID2TaskID)
	if !c.mapFinished {
		if state, ok := c.WorkerID2TaskID[id]; ok {
			state.lastHeatbeatTime = time.Now()
			c.WorkerID2TaskID[id] = state
		}
		for wid, state := range c.WorkerID2TaskID {
			if time.Since(state.lastHeatbeatTime).Seconds() > 10 {
				// fmt.Println(state.lastHeatbeatTime)
				// fmt.Printf("worker:%d taskid: %d 大于十秒。。。。。。。。。。。。。。。。。。。\n", wid, state.TaskID)
				tid := state.TaskID
				c.MapTaskList[tid] = Task{State: 0, WorkerID: 0}
				delete(c.WorkerID2TaskID, wid)
			}
		}
		return nil
	}

	if !c.Finished {
		if state, ok := c.WorkerID2TaskID1[id]; ok {
			state.lastHeatbeatTime = time.Now()
			c.WorkerID2TaskID1[id] = state
		}
		for wid, state := range c.WorkerID2TaskID1 {
			if time.Since(state.lastHeatbeatTime).Seconds() > 10 {
				// fmt.Println(state.lastHeatbeatTime)
				// fmt.Printf("+++worker:%d taskid: %d 大于十秒。。。。。。。。。。。。。。。。。。。\n", wid, state.TaskID)
				tid := state.TaskID
				c.ReduceTaskList[tid] = Task{State: 0, WorkerID: 0}
				delete(c.WorkerID2TaskID1, wid)
			}
		}
		return nil
	}

	return nil
}

func (c *Coordinator) GetWorkerID(req *GetWorkIDReq, resp *GetWorkIDResp) error {

	resp.WorkerID = c.GenerateWorkerID()
	return nil
}

// start a thread that listens for RPCs from worker.go
func (c *Coordinator) server() {
	rpc.Register(c)
	rpc.HandleHTTP()
	//l, e := net.Listen("tcp", ":1234")
	sockname := coordinatorSock()
	os.Remove(sockname)
	l, e := net.Listen("unix", sockname)
	if e != nil {
		log.Fatal("listen error:", e)
	}
	go http.Serve(l, nil)
}

// main/mrcoordinator.go calls Done() periodically to find out
// if the entire job has finished.
func (c *Coordinator) Done() bool {
	ret := false

	// Your code here.
	//所有任务完成后 返回true
	c.mu.Lock()
	ret = c.Finished
	c.mu.Unlock()
	return ret
}

// create a Coordinator.
// main/mrcoordinator.go calls this function.
// nReduce is the number of reduce tasks to use.
func MakeCoordinator(files []string, nReduce int) *Coordinator {
	mapTaskNum := len(files)
	c := Coordinator{
		nReduce:          nReduce,
		mapFinished:      false,
		Finished:         false,
		WorkerID2TaskID:  map[int]WorkerState{},
		WorkerID2TaskID1: map[int]WorkerState{},
		MapTaskList:      make([]Task, mapTaskNum),
		ReduceTaskList:   make([]Task, nReduce),
		mu:               sync.Mutex{},
		files:            files,
	}

	// Your code here.

	// 将待处理的文件列表进行数据分片

	// distributeWords(files, nReduce, "tmp")

	c.server()
	return &c
}

// distributeWords 将多个输入文件中的单词分发到 n 个输出文件中
func distributeWords(inputFiles []string, n int, outputDir string) error {
	// 创建输出目录
	if err := os.MkdirAll(outputDir, os.ModePerm); err != nil {
		return fmt.Errorf("could not create output directory: %w", err)
	}
	// 创建 n 个输出文件的句柄
	outputFiles := make([]*os.File, n)
	for i := 0; i < n; i++ {

		filePath := filepath.Join(outputDir, fmt.Sprintf("output_part_%d.txt", i+1))
		file, err := os.Create(filePath)
		if err != nil {
			return fmt.Errorf("could not create output file %s: %w", filePath, err)
		}
		outputFiles[i] = file
		defer file.Close() // 程序结束时关闭文件
	}

	// 按顺序读取输入文件中的单词并分发到输出文件
	currentFileIndex := 0
	for _, inputFile := range inputFiles {
		file, err := os.Open(inputFile)

		if err != nil {
			return fmt.Errorf("could not open input file %s: %w", inputFile, err)
		}
		defer file.Close()

		scanner := bufio.NewScanner(file)
		for scanner.Scan() {
			word := scanner.Text()
			_, err := outputFiles[currentFileIndex].WriteString(word + "\n")
			if err != nil {
				return fmt.Errorf("could not write to output file: %w", err)
			}

			// 更新当前文件索引，使单词轮流写入不同的文件
			currentFileIndex = (currentFileIndex + 1) % n
		}

		if err := scanner.Err(); err != nil {
			return fmt.Errorf("error reading file %s: %w", inputFile, err)
		}
	}

	return nil
}
